feat: add IndexTransform library for composable, lazy coordinate mappings - #3906
feat: add IndexTransform library for composable, lazy coordinate mappings#3906d-v-b wants to merge 91 commits into
Conversation
…ings Add a new `src/zarr/core/transforms/` package implementing TensorStore-inspired index transforms. The core idea: every indexing operation (slicing, fancy indexing, etc.) produces a coordinate mapping from user space to storage space. These mappings compose lazily — no I/O until explicitly resolved. Key types: - `IndexDomain` — rectangular region in N-dimensional integer space - `ConstantMap`, `DimensionMap`, `ArrayMap` — three representations of a set of storage coordinates (singleton, arithmetic progression, explicit enumeration) - `IndexTransform` — pairs an input domain with output maps (one per storage dim) - `compose(outer, inner)` — chain two transforms Key operations on IndexTransform: - `__getitem__`, `.oindex[]`, `.vindex[]` — indexing produces new transforms - `.intersect(domain)` — restrict to coordinates within a region (chunk resolution) - `.translate(shift)` — shift coordinates (make chunk-local) The transform library is standalone with no dependency on Array. Includes comprehensive test suite (143 tests covering all types, operations, composition, chunk resolution, and edge cases). Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #3906 +/- ##
==========================================
- Coverage 93.50% 92.90% -0.61%
==========================================
Files 90 98 +8
Lines 11981 13250 +1269
==========================================
+ Hits 11203 12310 +1107
- Misses 778 940 +162
🚀 New features to boost your workflow:
|
|
@d-v-b I'm new to zarr-python indexing, does My use case is mainly if I have an array of shape |
Add TypedDict definitions and conversion functions for serializing
IndexDomain, OutputIndexMap, and IndexTransform to/from JSON.
The JSON format follows TensorStore's conventions for interoperability:
- IndexDomain: input_inclusive_min, input_exclusive_max, input_labels
- OutputIndexMap: offset + optional stride/input_dimension/index_array
- IndexTransform: domain fields + output array
TypedDicts: IndexDomainJSON, OutputIndexMapJSON, IndexTransformJSON
Functions: index_domain_to_json, index_domain_from_json,
index_transform_to_json, index_transform_from_json
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
on this branch: >>> arr = create_array(store={}, shape=(100,100), dtype="uint8")
>>> arr.z[0]
<Array memory://136980013828096 shape=(100,) dtype=uint8 domain={ 0, [0, 100) }>
>>> arr.z[0].shape
(100,)the |
|
What is I have and I want to select the arrays for band 1 (second dim) but for all the times, would |
yeah it should! |
Merging this PR will degrade performance by 17.94%
Performance ChangesComparing Footnotes
|
|
I played with this a bit today and this type of slicing is really great: One thing that tripped me up -- probably just shows my naivety (PBCAK) -- is that e.g. data.z.shape is |
We should probably have a shape attribute there. Thanks for reporting that, and thanks for trying this branch out. Let me know if you find anything else we need to fix! |
|
I guess from my PoV it would be nice if I could just use Edit: Don't get me wrong, this is pretty great -- as an example, with this I can remove the dask.array wrapper here: instead, just check if my array is zarr and if so use |
I'm glad it's useful! I agree that this is probably how slicing should work by default. But that would be a big breaking change, so it's far down the road. The |
|
Makes sense. So maybe in followup _LazyIndexAccessor could have more of the array properties and support more of the numpy protocol, so that array.z could be used as a drop-in for array, but be lazy about slicing, etc. |
|
I like this a lot! |
I just wanted something really short. Since this is (IMO) a strictly better slicing API, I wanted to make accessing it as low-friction as possible. But I will totally bow to the will of the crowd here, we can use whatever people find intuitive.
I use |
|
I agree with @normanrz, the the length of |
|
what about |
|
Just for reference: xarray uses |
|
ping @ilan-gold for visibility |
ilan-gold
left a comment
There was a problem hiding this comment.
use data structures compatible with tensorstore. tensorstore's indexing machinery serializes to JSON. if we adopt the same patterns, we are closer to using tensorstore as an optional backend.
I'm sure zarrs could also support this but I wonder about JSON performance for things like vindex/oindex where you could easily have 10's of thousands of individual coordinates.
| - **translate(shift)** — shift all output coordinates. This makes coordinates | ||
| chunk-local: "express my coordinates relative to the chunk origin." | ||
|
|
||
| - **compose(outer, inner)** — chain two transforms. See ``composition.py``. |
There was a problem hiding this comment.
I think it could be cool to have different kinds of composition like union and I don't see that anywhere intthe PR - it seems that only array.z[my_transform][my_other_transform_relative_to_the_first] is supported wheres it could be cool to have array.z[my_transform + my_other_transform]. Concretely, if I wanted (slice(0, 10), slice(0, 10)) and then also (slice(50, 60), slice(50, 60)) from an array, how would I do that? np.concatenate + np.arange to generate outer indexers is the only option AFAICT but that kind of stinks.
There was a problem hiding this comment.
that's a very cool idea, not sure about overloading the addition operator, but some method for combining indexing expressions seems nice. I'll see what we need for this
…forms package Move the TensorStore-style index-transform algebra out of zarr core into a new numpy-only uv-workspace subpackage packages/zarr-transforms (import name zarr_transforms), mirroring packages/zarr-metadata. This is a pure move: behavior is byte-for-byte unchanged and the full suite stays green. The package depends on numpy + stdlib only and must not import zarr. Two couplings are broken: - errors: the canonical BoundsCheckError / VindexInvalidSelectionError class definitions now live in zarr_transforms.errors (both still subclass IndexError). zarr.errors re-exports the same objects, so zarr.errors.BoundsCheckError is zarr_transforms.errors.BoundsCheckError and every existing catch site is unaffected. - chunk grid: chunk_resolution is typed against new structural Protocols (ChunkGridLike / DimensionGridLike) in zarr_transforms.grid instead of importing zarr's concrete ChunkGrid, which satisfies them structurally. zarr now declares a runtime dependency on zarr-transforms, resolved to the in-tree package via a uv workspace source. The package tests are collected by the root test suite (they exercise chunk resolution against zarr's ChunkGrid, so they need zarr importable) rather than run in isolation. The package __init__ promotes the names the zarr integration layer consumes (iter_chunk_transforms, sub_transform_to_selections, selection_to_transform) to the public surface alongside the existing exports. Assisted-by: ClaudeCode:claude-opus-4.8
Assisted-by: ClaudeCode:claude-opus-4.8
The hatch run-coverage / run-coverage-html / run-hypothesis (and gputest run-coverage) scripts passed --source=src to coverage run, which excluded the index-transform algebra after its move to packages/zarr-transforms/src — it was measured before the move. Declare the measured source trees in [tool.coverage.run] source instead, so every coverage run invocation measures both src and the workspace package consistently, and drop the CLI flag from all four scripts. Assisted-by: ClaudeCode:claude-opus-4.8
Great idea, and thanks for offering your help. I'm working on moving the lazy indexing logic here out into a separate package in the zarr-python workspace. That PR should be much easier to land, since it wouldn't change anything about zarr python. Once that package is up and running, we can publish it on pypi and you can start trying it out. How does this sound?
I'm happy to prototype either tier if that's a useful thing to peel off. I think we can sidestep this for now, because the JSON thing is only relevant when / if we want to interchange array selections expressions over the wire. We don't need a performant JSON serialization inside Zarr-Python itself.
this is a shared need -- see #4028. I think the lazy-indexing-oriented API would look like this: you declare the N regions of the array you want to read, then you arrange those N regions into a concatenated array, and then you submit the request to fetch that array to a zarr backend that can efficiently do all the IO you need. This is just an idea right now, but I think it would be a very clean API. |
Applying an oindex/vindex selection to a view that already carries an
orthogonal ArrayMap could land a fancy index on a broadcast (singleton)
axis of that map, leaking a raw numpy IndexError ("index N is out of
bounds for axis ... with size 1") at resolve time — reading like a user
error rather than the implementation gap it is.
Add `_guard_fancy_after_fancy`, invoked from `_apply_oindex` and
`_apply_vindex`, which raises a clear NotImplementedError at composition
time when a fancy axis targets a non-dependency axis of an existing
ArrayMap. It names the limitation and the workaround (materialize via
.result() then index, or reorder so the fancy step is last). Compositions
that keep the fancy step on the dependency axis (or on a correlated vindex
view) are unaffected and still resolve correctly.
Adds TestFancyAfterFancy, documents the limitation in the lazy-indexing
user guide, and refreshes the now-stale generator/runner comments that
described this class as a silent bug.
Assisted-by: ClaudeCode:claude-fable-5
…anges
zarr-transforms is a hard runtime dependency of zarr with no CI. Add
`zarr-transforms.yml` (push/PR path-filtered test/ruff/pyright, mirroring
zarr-metadata; the test job runs from the repo root because the transform
tests import zarr) and `zarr-transforms-release.yml` (tag-triggered publish
on `zarr_transforms-v*`). Extend `check_changelogs.yml` to check the
package's changes directory.
Bump the package `requires-python` floor to >=3.12 (consistent with the
repo; nothing tested 3.11) and update its classifiers, ruff target, and
pyright pythonVersion to match. Extend the root hatch `--match v*` comment
to note the `zarr_transforms-v*` tags it also excludes.
Disclose the two deliberate eager-path changes in the 3906 changelog: the
Array repr `domain={...}` suffix and the 0-d iteration TypeError matching
NumPy. Fix a stale comment in array.py that claimed pop_fields yields []
for no fields (it now yields None).
Assisted-by: ClaudeCode:claude-fable-5
Mirrors zarr_metadata's importlib.metadata idiom; the release workflow's isolated-wheel check imports it. Assisted-by: ClaudeCode:claude-fable-5
…ansforms The ArrayMap (fancy) branch of iter_chunk_transforms enumerated the dense range(min_chunk, max_chunk + 1) bounding box and ran transform.intersect against every candidate chunk, making sparse fancy/vindex chunk resolution scale O(n_chunks) instead of O(n_touched). The exact touched chunk ids were already computed and then discarded in favor of min/max. Enumerate each fancy dimension's distinct touched chunk ids (np.unique) instead; the cartesian product then spans only touched-per-dimension combinations. Constant/Dimension dims keep their contiguous ranges. Semantics are unchanged: the dense range only ever added empty intersections, which intersect already skipped. Sparse vindex (2 far-apart coords) goes from 15.9x slower than eager at 1k chunks / 65.5x at 4k to ~1.1x at both, flat in grid size. Dense fancy selections are unchanged. Assisted-by: ClaudeCode:claude-fable-5
Vendored, unmodified, from ndsel branch fix/slice-origin-trunc (commit c132b4c1caa3205830ce35a42502363171f650a7). PROVENANCE.md records the source URL, commit SHA, and do-not-edit note. The corpus is the language-agnostic fixture set every ndsel implementation runs against. Assisted-by: ClaudeCode:claude-opus-4.8
Add a pure JSON->JSON message layer (messages.py: parse_ndsel,
normalize_ndsel, NdselError) implementing the ndsel draft wire format,
which adapts TensorStore's IndexTransform. It accepts all five kinds
(point/box/slice/points/transform), normalizes to the canonical transform
body (spec 4.3), and enforces the full error taxonomy. Verified against the
vendored conformance corpus.
Rework json.py into the engine/lowering layer: transform_{to,from}_canonical
(re-pointed from index_transform_{to,from}_json). Engine constraints live
here only (reject infinite bounds, implicit-lowers-by-value). Fix the oindex
wire format so index_array maps no longer carry input_dimension (rejected by
ndsel and TensorStore); degenerate all-singleton arrays collapse to constant
maps, and input_dimension is reconstructed from array dependency axes on load.
Cross-checked against a real tensorstore (importorskip) for a set of finite
canonical bodies.
Assisted-by: ClaudeCode:claude-opus-4.8
The pyright job added alongside zarr-metadata's config had never actually run before an external PR triggered it, and fails with 68 errors on this branch's HEAD. Downgrade reportUnknownVariableType/Argument/Member/ParameterType to warnings: the strict config was copied from zarr-metadata's JSON/dataclass code, but numpy's stubs return partially-unknown types even at fully-typed call sites, so this numpy-heavy package can't reasonably satisfy that family. CI only fails on errors, so warnings keep the signal visible without blocking the build. Fix the genuine findings in code: - reportUnnecessaryIsInstance (14 sites across transform.py, composition.py, json.py, chunk_resolution.py): replace tautological final isinstance arms (pyright proves them always-true once prior branches exhaust the union) with `else:` plus a comment naming the narrowed type. Two sites in composition.py and one in json.py also drop a `# pragma: no cover`-marked dead `raise TypeError` that followed. - reportPrivateUsage (2 sites): chunk_resolution.py's `chunk_grid._dimensions` (declared on the ChunkGridLike Protocol for structural typing against zarr's ChunkGrid) and json.py's cross-module import of transform.py's `_array_map_dependency_axes`. Both are suppressed with `# pyright: ignore[reportPrivateUsage]` plus a comment; renaming either is left as an open pre-publish API decision. Zero behavior change. Verified: pyright 0 errors (was 68); pytest 266 passed, 1 skipped; mypy clean; ruff clean. Assisted-by: ClaudeCode:claude-fable-5
Fixtures verified byte-identical to the merged c59bc556c; only the reference moves off the deleted feature branch. Assisted-by: ClaudeCode:claude-fable-5
…nsforms The fast path bound 'm' as ArrayMap before the general loop rebinds it to the OutputIndexMap union, which mypy rejects (and treats the ConstantMap arm as unreachable). packages/ is outside the pre-commit mypy scope, so #234's CI could not catch this. Assisted-by: ClaudeCode:claude-fable-5
…ndexing # Conflicts: # pyproject.toml # src/zarr/core/array.py # uv.lock
…indexing Pre-publish rename: the in-tree zarr-transforms workspace package is not yet published, so this is a safe, purely mechanical rename with zero behavior change. Distribution name is now zarr-indexing; import name is now zarr_indexing. Covers: the package directory and src layout, the package's own pyproject.toml (name, tag-pattern, wheel packages, towncrier package), the root pyproject.toml (workspace member, uv source, runtime dependency, coverage source, pytest testpaths, hatch version comment), all imports across src/zarr, tests, and the package's own tests, the zarr-transforms(-release).yml workflows (renamed to zarr-indexing(-release).yml with matching job names, path filters, tag pattern, environment names, artifact names, and smoke test), check_changelogs.yml's package path, changelog fragments (root changes/3906.feature.md and the package's own changes/), and the regenerated uv.lock. Assisted-by: ClaudeCode:claude-fable-5
…classes
Identity (non-view) arrays now render with the legacy 3.x repr with no
`domain={...}` suffix, byte-identical across `Array` and `AsyncArray`.
The suffix appears only on non-identity lazy views, and `AsyncArray`
gains it for views (reachable via `translate_by` and friends). Restores
the doctests and repr tests that had been rewritten to expect the
unconditional suffix on identity arrays.
Assisted-by: ClaudeCode:claude-opus-4.8
`get_coordinate_selection(out=...)` requires a flat `(n,)` buffer on eager/identity arrays but the broadcast selection shape on lazy views. Document both sides in the `out` parameter docstring and note that the broadcast-shape contract is the long-term direction. No behavior change. Assisted-by: ClaudeCode:claude-opus-4.8
…nning docs) - name the zarr-indexing changelog fragment by PR number (3906) so the integer-issue check passes - drop :class: sphinx roles from the indexing-program docstrings (mkdocs markdown; caught by main's new ci/lint_docs.py) - untrack docs/superpowers/ planning documents: upstream main gitignores that path (agent planning notes are local-only by convention), which also removes them from markdownlint's scope; files remain on disk Assisted-by: ClaudeCode:claude-fable-5
…dLike Protocol Pre-publish API cleanup (option 2): iter_chunk_transforms now accepts Sequence[DimensionGridLike] directly, so the published contract contains only public names. zarr's integration layer passes its own private chunk_grid._dimensions, keeping the private access inside zarr where it belongs; both pyright reportPrivateUsage suppressions tied to the old Protocol are gone. Assisted-by: ClaudeCode:claude-fable-5
The copy predates the merge from main that bumped setup-uv to v8.3.2 and attest to v4.1.1 in the zarr-metadata release workflow. Assisted-by: ClaudeCode:claude-fable-5
5990d89 to
43efa85
Compare
Assisted-by: ClaudeCode:claude-fable-5
…y paths `Order.check` differenced adjacent elements with `np.diff`, so on an unsigned dtype a descending step wrapped around to a large positive value and the selection was classified as increasing. `IntArrayDimIndexer` then took the no-reorder fast path and read the wrong region out of each chunk. Unsigned arrays also reached the chunk-relative arithmetic, where NumPy promotes `uint64` mixed with a signed offset to `float64`, yielding a float chunk selection and `IndexError: arrays used as indices must be of integer (or boolean) type`. `Order.check` now compares adjacent elements instead of differencing them, and both eager entry points (`IntArrayDimIndexer` for oindex, `CoordinateIndexer` for vindex) cast unsigned index arrays to `np.intp` before any order check or chunk arithmetic. A value too large for `np.intp` cannot index any array, so it raises `BoundsCheckError` rather than wrapping to a negative index. Signed input is untouched. The lazy path already normalized to `np.intp` at its boundary and is unaffected. Assisted-by: ClaudeCode:claude-fable-5
My summary:
With dask maintenance on the decline, it's more important than ever that we give zarr-python users a dask-free way to do something very intuitive: index large zarr arrays without turning the whole thing into a numpy array first. This was discussed at length in #1603.
This PR, done with Claude, makes regular indexing go through a lazy indexing layer. The lazy indexing layer is based on abstractions defined in tensorstore. The basic idea is to explicitly model indexing an array as a transformation from some input coordinates to output coordinates, and to bind such a representation to our
Arrayclasses.Regular indexing via
.__getitem__is still immediate, but arrays have a new.zattribute that exposes the lazy indexing layer:Goals here:
mainare disparate ad-hoc copies of stuff from zarr-python 2.x. We can do better.Non-goals:
Claude's summary.
Add a new
src/zarr/core/transforms/package implementing TensorStore-inspiredindex transforms. The core idea: every indexing operation (slicing, fancy indexing,
etc.) produces a coordinate mapping from user space to storage space. These mappings
compose lazily — no I/O until explicitly resolved.
Key types:
IndexDomain— rectangular region in N-dimensional integer spaceConstantMap,DimensionMap,ArrayMap— three representations of a set ofstorage coordinates (singleton, arithmetic progression, explicit enumeration)
IndexTransform— pairs an input domain with output maps (one per storage dim)compose(outer, inner)— chain two transformsKey operations on IndexTransform:
__getitem__,.oindex[],.vindex[]— indexing produces new transforms.intersect(domain)— restrict to coordinates within a region (chunk resolution).translate(shift)— shift coordinates (make chunk-local)The transform library is standalone with no dependency on Array.
Includes comprehensive test suite (143 tests covering all types, operations,
composition, chunk resolution, and edge cases).
Co-Authored-By: Claude Opus 4.6 (1M context) noreply@anthropic.com